//////////////////////////////////////////////////////////////////
//                        COPYRIGHT NOTICE
// (c) Copyright 1998-2002 by Jeongnim Kim, Lars Jonsson and Seungwon Lee
//   All rights reserved.
//////////////////////////////////////////////////////////////////

//////////////////////////////////////////////////////////////////
//                    CONTACT INFORMATION
//   Jeongnim Kim
//   National Center for Supercomputing Applications &
//   Materials Computation Center
//   University of Illinois, Urbana-Champaign
//   Urbana, IL 61801
//   e-mail: jnkim@ncsa.uiuc.edu
//   Tel:    217-244-6319 (NCSA) 217-333-3324 (MCC)
//
// Supported by 
//   National Center for Supercomputing Applications, UIUC
//   Materials Computation Center, UIUC
//   Department of Physics, Ohio State University
//   Ohio Supercomputer Center
//////////////////////////////////////////////////////////////////
// -*- C++ -*-

#ifndef OHMMS_GENERIC_ARPACKSYMENGINE_H
#define OHMMS_GENERIC_ARPACKSYMENGINE_H

#include <complex>
#include <iostream>
#include <map>
#include <algorithm>
using namespace std;

#include "EigSolver/EigenSpaceDimension.h"
#include "EigSolver/Arpackf.h"
#include "OhmmsNumeric/OhmmsBlasf.h"
#ifdef OHMMS_MTL
#include "mtl/mtl.h"
#else
#include "OhmmsPETE/OhmmsVector.h"
#endif

#if defined(ADD_)
#define dpotrf  dpotrf_
#define dpotrs  dpotrs_
#endif
extern "C" {
  void dposv(char* uplo,int *,int *, double*, int*, double*, int*, int*);
  void dpotrf(char *uplo, int*, double*, int*, int*);
  void dpotrs(char *uplo, int*, int *, double*, int*, double*, int*, int*);

}
template<class T, class EigVec_t>
struct ArpackSymSolver {

  typedef typename OHMMS::type_traits<T>::value_type value_type;
  typedef typename OHMMS::type_traits<T>::real_type  real_type;
  typedef EigVec_t                 EigenStates_t;
#ifdef OHMMS_MTL
  typedef mtl::dense1D<value_type> WorkArray_t;
  typedef mtl::dense1D<real_type>  RealArray_t;
  typedef mtl::dense1D<int>        IndexArray_t;
#else
  typedef Vector<value_type> WorkArray_t;
  typedef Vector<real_type>  RealArray_t;
  typedef Vector<int>        IndexArray_t;
#endif

  typedef typename WorkArray_t::size_type   size_type;

  int N;
  int M;
  int ncv; // number of Lanczos vectors
  int info;
  int lworkl;
  int ido;
  int ierr;
  int nconv;
  int maxiter;
  int ishfts;
  int mode;
  int rvec;
  string which;
  int iparam[11], ipntr[14];
  char bmat;

  IndexArray_t select;
  WorkArray_t workd; // array of 3N+1
  WorkArray_t workl; // array of lworkl
  WorkArray_t resid; // residue
  WorkArray_t eigVal;// eigen values
  WorkArray_t eigVec;// eigen vectors 
  RealArray_t rwork; // only for a complex problem

  real_type tol;  // tolerance
  value_type sigma; // reference value for shift-invert


  ArpackSymSolver(EigenSpaceDimension& dim, bool nonortho=false){   
    N = 0;
    M = 0;
    tol = 1.0e-13;
    info = 0;
    ido = 0;
    maxiter = 1000;
    ncv = -1;
    iparam[0] = 1; iparam[6] = 1;
    sigma = 0.0;
    which = "SM";
  }

  ~ArpackSymSolver() { }

  template<class HamiltonianMatrix_t>
  bool solve(HamiltonianMatrix_t& H) {

    bmat = 'I';
    initializeWorkSpace();
    iparam[0] = 1; iparam[6] = 1;
    ido = 0;
    do{
      solve(eigVec.data());
      H.apply(workd.data()+ipntr[0]-1, workd.data()+ipntr[1]-1);
    } while(ido == -1 || ido == 1);
    rvec = 1;
    finalize(eigVec.data());

    //Psi.resize(N,M);
    //OhmmsBlas::copy(N*M,eigVec.begin(),Psi.data());
    //std::copy(eigVec.begin(), eigVec.end(), Psi.data());
    if(info < 0) {
      cerr << endl << "Error with __aupd, info = " <<  info << endl;
      cerr << " Check documentation in __aupd " << endl << endl;
      return false;
    } else
      return true;
  }

  template<class HamiltonianMatrix_t>
  bool solve(HamiltonianMatrix_t& H, EigenStates_t& Psi) {
    bool done = solve(H);
    // copy eigen states according to the ascending eigenvalues
    multimap<double,int> ordered_e;
    for(int i=0; i<M; i++) ordered_e.insert(pair<double,int>(eigVal[i],i));
    multimap<double,int>::iterator it = ordered_e.begin();
    int ieig = 0;
    for(; it!=ordered_e.end(); it++) {
      eigVal[ieig] = (*it).first;      
      int is = (*it).second*N;
      std::copy(eigVec.data()+is,eigVec.data()+is+N, Psi.data()+ieig*N);
      ieig++; 
    }
    return done;
  }

  template<class HamiltonianMatrix_t>
  bool solve(HamiltonianMatrix_t& H, const HamiltonianMatrix_t& S) {

    bmat = 'G';
    char NoTranspose = 'N';

    initializeWorkSpace();

    EigenStates_t h(N,N);
    ido = 0;
    int one = 1;
    bool tocontinue = true;
    char uplo = 'L';
    int counter = 0;

    ///compute Cholesky factorization A  U'*U or A = L*L'

    iparam[6] = 2;
    OHMMS::copy(S,h);
    dpotrf(&uplo,&N,h.data(),&N,&info);
    
    do{

      solve(eigVec.data());
      
      if(ido == -1 || ido == 1) {
	
	value_type *work0 = workd.data()+ipntr[0]-1;
	value_type *work1 = workd.data()+ipntr[1]-1;
	
	H.apply(work0,work1);
	std::copy(work1,work1+N,work0);

	///solve the linear system A*X = B
	dpotrs(&uplo,&N,&one,h.data(),&N,work1,&N,&info);
      } else if(ido == 2) {
	value_type *work1 = workd.data()+ipntr[1]-1;
	S.apply(workd.data()+ipntr[0]-1, work1);
	//S.apply(workd.data()+ipntr[0]-1, workd.data()+ipntr[1]-1);
      } else {
	tocontinue = false;
      }
    } while(tocontinue);
    rvec = 1;
    finalize(eigVec.data());

    if(info < 0) {
      cout << endl << "Error with __aupd, info = " <<  info << endl;
      cout << " Check documentation in __aupd " << endl << endl;
      return false;
    } else
      return true;
//     } else if(which == "SM") {
//       ///using shift-inversion method
//       which = "LM";
//       iparam[0] = 1;
//       iparam[6] = 3;
//       /// h = H - sigma*S
//       //OHMMS::copy(H,S,-sigma,h);
//       OHMMS::copy(H,S,-sigma,h);
//       dpotrf(&uplo,&N,h.data(),&N,&info);
//       do{
// 	solve(eigVec.data());
// 	if(ido == -1) {
// 	  value_type *work0 = workd.data()+ipntr[0]-1;
// 	  value_type *work1 = workd.data()+ipntr[1]-1;
// 	  S.apply(work0,work1);
// 	  ///solve the linear system A*X = B
// 	  dpotrs(&uplo,&N,&one,h.data(),&N,work1,&N,&info);
// 	} else if(ido == 1) {
// 	  value_type *work2 = workd.data()+ipntr[2]-1;
// 	  value_type *work1 = workd.data()+ipntr[1]-1;
// 	  std::copy(work2,work2+N, work1);
// 	  ///solve the linear system A*X = B
// 	  dpotrs(&uplo,&N,&one,h.data(),&N,work1,&N,&info);
// 	} else if(ido == 2) {
// 	  S.apply(workd.data()+ipntr[0]-1, workd.data()+ipntr[1]-1);
// 	} else if(ido == 3) {
// 	  cout << "Shift needs to be evaluated." << endl;	   
// 	} else {
// 	  tocontinue = false;
// 	}
//       } while(tocontinue);
//     }
//     rvec = 1;
//     finalize(eigVec.data());
//     // copy eigen states according to the ascending eigenvalues
//     multimap<double,int> ordered_e;
//     for(int i=0; i<M; i++) ordered_e.insert(pair<double,int>(eigVal[i],i));
//     multimap<double,int>::iterator it = ordered_e.begin();
//     int ieig = 0;
//     for(; it!=ordered_e.end(); it++) {
//       eigVal[ieig] = (*it).first;      
//       int is = (*it).second*N;
//       std::copy(eigVec.data()+is,eigVec.data()+is+N, Psi.data()+ieig*N);
//       ieig++; 
//     }
//     if(info < 0) {
//       cout << endl << "Error with __aupd, info = " <<  info << endl;
//       cout << " Check documentation in __aupd " << endl << endl;
//       return false;
//     } else
//       return true;
  }

  template<class HamiltonianMatrix_t>
  bool solve(HamiltonianMatrix_t& H, const HamiltonianMatrix_t& S, 
	     EigenStates_t& Psi) {
    bool done = solve(H,S);
    std::copy(eigVec.data(),eigVec.data()+M*N, Psi.data());
    return done;
  }

  template<class HamiltonianMatrix_t>
  bool checkEigvectors(HamiltonianMatrix_t& H, ostream& os) {
    if ( ierr != 0) {
      os << endl << "Error with __eupd, info = " << ierr << endl;
      os << "Check the documentation of __eupd. " << endl;
      return false;
    } else {
      
      nconv =  iparam[4];
      if(rvec == 1) {
	WorkArray_t ax(N);
	WorkArray_t res(M);
	
	for(int j=0; j<nconv; j++) { // residual calculations
	  ax = 0.0;
	  H.apply(eigVec.data()+j*N, ax.data());
	  value_type da  = -eigVal[j];
	  int incx = 1;
	  int incy = 1;
	  OhmmsBlas::axpy(N, da, eigVec.data()+j*N, incx, ax.data(), incy);  
	  res[j] = OhmmsBlas::norm2(N, ax.data(), incx);
	  res[j] /= fabs(eigVal[j]);
	}
	os << "Ritz values and relative residuals " << endl;
	for(int j=0; j<nconv; j++)
	  os << j << "  "  << eigVal[j] << "  " << res[j] << endl;
      } else {
	os << "Ritz values only " << endl;
	for(int j=0; j<nconv; j++)
	  os << j << "  "  << eigVal[j] <<  endl;
      }

      os << endl;
      if(info == 1) {
	os << endl << " Maximum number of iterations reached." << endl;
      } else if ( info == 3) {
	os << endl << "No shifts could be applied during implicit "
	     <<   "  Arnoldi update, try increasing NCV." << endl;
      }
      os << endl;
      os << " _SDRV1 " << endl;
      os << " ====== " << endl << endl;
      os << " Size of the matrix is " << N << endl;
      os << " The number of Ritz values requested is " << M << endl;
      os << " The number of Arnoldi vectors generated  (NCV) is " << ncv << endl;
      os << " What portion of the spectrum: " << which << endl;
      os << " The number of converged Ritz values is " << nconv << endl;
      
      os << " The number of Implicit Arnoldi update iterations taken is ";
      os <<  iparam[2] << endl;
      os << " The number of OP*x is " << iparam[8] << endl;
      os << " The convergence criterion is " << tol << endl;
      return true;
    }
  }

  template<class HamiltonianMatrix_t>
  bool checkEigvectors(const HamiltonianMatrix_t& H, 
		       const HamiltonianMatrix_t& S,
		       ostream& os) {
    if ( ierr != 0) {
      os << endl << "Error with __eupd, info = " << ierr << endl;
      os << "Check the documentation of __eupd. " << endl;
      return false;
    } else {
      
      nconv =  iparam[4];
      if(rvec == 1) {
	WorkArray_t ax(N),mx(N);
	WorkArray_t res(M);
	
	for(int j=0; j<nconv; j++) { // residual calculations
	  ax = 0.0;
	  mx = 0.0;
	  H.apply(eigVec.data()+j*N, ax.data());
	  S.apply(eigVec.data()+j*N, mx.data());
	  value_type da  = -eigVal[j];
	  int incx = 1;
	  int incy = 1;
	  OhmmsBlas::axpy(N, da, mx.data(), incx, ax.data(), incy);  
	  res[j] = OhmmsBlas::norm2(N, ax.data(), incx);
	  res[j] /= fabs(eigVal[j]);
	}
	os << "Ritz values and relative residuals (shifted by " << sigma << " ) " <<endl;
	for(int j=0; j<nconv; j++)
	  os << j << "  "  << eigVal[j]+sigma << "  " << res[j] << endl;
      } else {
	os << "Ritz values only " << endl;
	for(int j=0; j<nconv; j++)
	  os << j << "  "  << eigVal[j]+sigma <<  endl;
      }

      os << endl;
      if(info == 1) {
	os << endl << " Maximum number of iterations reached." << endl;
      } else if ( info == 3) {
	os << endl << "No shifts could be applied during implicit "
	     <<   "  Arnoldi update, try increasing NCV." << endl;
      }
      os << endl;
      os << " _SDRV1 " << endl;
      os << " ====== " << endl << endl;
      os << " Size of the matrix is " << N << endl;
      os << " The number of Ritz values requested is " << M << endl;
      os << " The number of Arnoldi vectors generated  (NCV) is " << ncv << endl;
      os << " What portion of the spectrum: " << which << endl;
      os << " The number of converged Ritz values is " << nconv << endl;
      
      os << " The number of Implicit Arnoldi update iterations taken is ";
      os <<  iparam[2] << endl;
      os << " The number of OP*x is " << iparam[8] << endl;
      os << " The convergence criterion is " << tol << endl;
      return true;
    }
  }

  inline void setMaxIter(int n) { maxiter = n;}
  inline void setMaxTrialStates(int n) { ncv = n;}
  inline void setTolerance(real_type x) { tol = x;}
  inline real_type eigval(size_type i) const { return eigVal[i];}

  //bool put(xmlDocPtr doc, xmlNsPtr ns, xmlNodePtr cur);
  bool put(xmlDocPtr doc,xmlNodeSetPtr nodeset);

  void initializeWorkSpace() {

    if(ncv < 0) ncv = 4*M;
    if(ncv > N) ncv = N;
    //TODO: COMPLEX and DOUBLE  (maybe not)
    lworkl = ncv*(3*ncv+5);
    iparam[2] = maxiter;
    if(workd.size() < 3*N) {
      workd.resize(3*N);
    }                                    
    if(resid.size() < N) {
      resid.resize(N);
    }
    if(select.size() < ncv) {
      select.resize(ncv);
    }
    if(eigVal.size() < ncv*2) {
      eigVal.resize(ncv*2);
    }
    if(eigVec.size() < ncv*N) {
      eigVec.resize(ncv*N);
    }
    if(workl.size() < lworkl) {
      workl.resize(lworkl);
    }
    
  }

  void solve(double* eigv){
    int n = N; int ldv = N; int nev = M;
    dsaupd(&ido, &bmat, &n, which.c_str(), &nev, &tol,
           resid.data(), &ncv, eigv, &ldv, iparam, ipntr,
           workd.data(), workl.data(), &lworkl, &info);
  }

  void solve(complex<double>* eigv){

    if(rwork.empty()) {
      rwork = RArray_t(ncv); // only for a complex problem
    } 
    int n = N; int ldv = N; int nev = M;
    znaupd(&ido, &bmat, &n, which.c_str(), &nev, &tol,
           resid.data(), &ncv, eigv, &ldv, iparam, ipntr,
           workd.data(), workl.data(), &lworkl, rwork.data(), &info);
  }

  void solve(float* eigv){
    int n = N; int ldv = N; int nev = M;
    ssaupd(&ido, &bmat, &n, which.c_str(), &nev, &tol,
           resid.data(), &ncv, eigv, &ldv, iparam, ipntr,
           workd.data(), workl.data(), &lworkl, &info);
  }

  void finalize(double *eigv) {
    char howmany ='A';
    int n = N; int ldv = N; int nev = M;
    dseupd( &rvec, &howmany, select.data(),
            eigVal.data(), eigv, &ldv, &sigma,
            &bmat, &n, which.c_str(), &nev, &tol, resid.data(), &ncv,
            eigv, &ldv, iparam, ipntr,
            workd.data(), workl.data(), &lworkl, &ierr );

    //for(int i=0; i<States.M; i++) ene[i] = eigVal[i];
  }

  void finalize(complex<double> *eigv) {
    char howmany ='A';
    int n = N; int ldv = N; int nev = M;
    WorkArray_t workev(2*ncv);    
    zneupd( &rvec, &howmany, select.data(),
            eigVal.data(), eigv, &ldv, &sigma, 
            workev.data(),  &bmat, &n, which.c_str(), &nev, &tol, 
            resid.data(), &ncv, eigv, &ldv, iparam, ipntr, workd.data(), 
            workl.data(), &lworkl, rwork.data(), &ierr );
  }

  void finalize(float *eigv) {
    char howmany ='A';
    int n = N; int ldv = N; int nev = M;
    sseupd( &rvec, &howmany, select.data(),
            eigVal.data(), eigv, &ldv, &sigma,
            &bmat, &n, which.c_str(), &nev, &tol, resid.data(), &ncv,
            eigv, &ldv, iparam, ipntr,
            workd.data(), workl.data(), &lworkl, &ierr );
  }
};
#endif

/***************************************************************************
 * $RCSfile: ArpackSymSolver.h,v $   $Author: jnkim $
 * $Revision: 1.1.1.1 $   $Date: 2004/07/23 21:04:02 $
 * $Id: ArpackSymSolver.h,v 1.1.1.1 2004/07/23 21:04:02 jnkim Exp $ 
 ***************************************************************************/

